
(* data types -----------------------------------------------------------------
*)


exception No_solution;

type state = (int * int);

type node = {
  state : state;
  path_cost : int;
  remaining_cost : int;
  total_cost : int;
  parent : option node
};

type operator_result = {
  next_state : state;
  cost : int
};


value string_of_node n = 
  Printf.sprintf 
    "{state=(%d,%d);path_cost=%d;remaining_cost=%d;total_cost=%d}" 
    (fst n.state) (snd n.state) n.path_cost n.remaining_cost n.total_cost;

module DHeap = struct (*{{{*)
  exception Empty;

  value nothing = {state=(0,0);path_cost=0;remaining_cost=0;total_cost=0;parent=None};

  value less_than a b = 
    a.total_cost < b.total_cost
    || 
    (a.total_cost == b.total_cost && a.path_cost > b.path_cost);

  value size = ref 0;
  value array : ref (array node) = ref [| |];
                                       
  value clear () = size.val := 0;

  value reallocate new_length =
    array.val := 
      Array.append 
        !array 
        (Array.make (new_length-Array.length !array) nothing);

  value push n =
  (
    if !size + 1 > (Array.length !array) 
    then reallocate ((!size+1)*2)
    else ();
    let rec move_up index =
      let parent_index = (index+1)/2-1
      in
        if index != 0 && less_than n (!array.(parent_index)) 
        then 
        (
          !array.(index) := !array.(parent_index);
          move_up parent_index
        )
        else
          !array.(index) := n
    in
    (
      move_up !size;
      incr size;
(*       size <- size + 1 *)
    );
  );
        
  value pop () =
    if !size = 0 
    then raise Empty
    else 
    (
      decr size;
(*       size <- size - 1; *)
      let result = !array.(0)
      in
        if !size = 0 
        then result
        else 
          let moving_value = !array.(!size) in
          let rec move_down tier_start next_tier_start index =
            let index_of_least = ref (-1)
            and least = ref moving_value
            and left_kid = next_tier_start+2*(index-tier_start) in
            let right_kid = left_kid + 1
            in
            (
              if (* left child in tree *) left_kid < !size then
                if less_than !array.(left_kid) !least 
                then
                (
                  index_of_least.val := left_kid;
                  least.val := !array.(left_kid)
                )
                else ()
              else ();
              if (* right child in tree *) right_kid < !size 
              then 
                if less_than !array.(right_kid) !least 
                then 
                (
                  index_of_least.val := right_kid;
                  least.val := !array.(right_kid)
                )
                else ()
              else ();
              if !index_of_least = -1 
              then
                (* heap condition satisfied *)
                !array.(index) := moving_value
              else 
              (
                (* heap condition not satisfied, recurse *)
                !array.(index) := !array.(!index_of_least);
                move_down next_tier_start ((2*(next_tier_start+1))-1) !index_of_least
              )
            )
          in
          (
            move_down 0 1 0;
            result
          );
    );

end;(*}}}*)

(* tool functions -------------------------------------------------------------
*)


value rec gobble_up_states_to_root ~node ~state_list =
  match node.parent with
  [ Some(parent) ->
      (gobble_up_states_to_root
         ~node:parent
         ~state_list:([ node.state :: state_list]))
  | None -> state_list
  ];

value used_states = HashSet.create 40;


(* A* search ------------------------------------------------------------------
*)


value a_star ~initial_state ~generate_successor_states ~estimate_remaining_cost =

  let inital_cost = estimate_remaining_cost initial_state in
  let initial_node = 
    {
      state = initial_state;
      path_cost = 0;
      remaining_cost = inital_cost;
      total_cost = inital_cost;
      parent = None
    }
  in
  let make_node ~parent op_result =
    let path_cost = op_result.cost + parent.path_cost
    and remaining_cost = estimate_remaining_cost op_result.next_state 
    in
      {
        state = op_result.next_state;
        path_cost = path_cost;
        remaining_cost = remaining_cost;
        total_cost = path_cost + remaining_cost;
        parent = Some(parent)
      } 
  in
  let () = DHeap.clear () in

  (*
  let node_queue = 
    new Heap.heap 
      ~less_than: (fun a b -> a.total_cost < b.total_cost
                     || (a.total_cost == b.total_cost 
                                           && a.path_cost > b.path_cost))
      ~nothing: initial_node
  in
  *)
  let goal_test node = node.remaining_cost = 0 in
  let rec inner_a_star () =
    let currently_expanded_node = DHeap.pop () in
(*     let () = Printf.printf "poped node: %s\n%!" (string_of_node currently_expanded_node) in *)
    if goal_test currently_expanded_node 
    then
      gobble_up_states_to_root ~node:currently_expanded_node ~state_list:[ ]
    else
     (
      List.iter (fun s ->
         if not (HashSet.mem used_states s.next_state)
         then
           let node = make_node ~parent:currently_expanded_node s in
             ( 
               HashSet.add used_states s.next_state;
               DHeap.push node;
             )
         else ()
      ) (generate_successor_states currently_expanded_node.state);
      inner_a_star ()
     )
  in
  (
    DHeap.push initial_node;
    HashSet.clear used_states;
    HashSet.add used_states initial_state;
    try
      inner_a_star ();
    with [ DHeap.Empty -> raise No_solution ];
  );
